"""
流体物性参数计算模块

该模块使用CoolProp库计算水和水蒸气在不同温度和压力下的物性参数。
包括:
1. 基本物性(密度、粘度、导热系数等)
2. 饱和状态物性
3. 传热相关参数(普朗特数、雷诺数等)
"""

import numpy as np
from typing import Dict, Tuple, Optional
from enum import Enum
from CoolProp.CoolProp import PropsSI

class FluidState(Enum):
    """流体状态枚举"""
    LIQUID = 1      # 液相
    VAPOR = 2       # 气相
    SATURATED = 3   # 饱和状态

class FluidProperties:
    """流体物性计算类"""
    
    def __init__(self):
        """初始化流体物性计算器"""
        self.fluid = "Water"  # 使用水作为工质
        
    def get_properties(self, pressure: float, temperature: float, 
                      state: FluidState) -> Dict[str, float]:
        """
        计算给定状态下的流体物性
        
        参数:
            pressure: 压力 (Pa)
            temperature: 温度 (K)
            state: 流体状态
            
        返回:
            包含以下物性的字典:
            - density: 密度 (kg/m³)
            - viscosity: 动力粘度 (Pa·s)
            - conductivity: 导热系数 (W/m·K)
            - specific_heat: 定压比热容 (J/kg·K)
            - enthalpy: 比焓 (J/kg)
            - surface_tension: 表面张力 (N/m) (仅饱和状态)
        """
        try:
            if state == FluidState.SATURATED:
                return self._get_saturation_properties(pressure)
                
            # 基本物性计算
            props = {
                'density': PropsSI('D', 'P', pressure, 'T', temperature, self.fluid),
                'viscosity': PropsSI('V', 'P', pressure, 'T', temperature, self.fluid),
                'conductivity': PropsSI('L', 'P', pressure, 'T', temperature, self.fluid),
                'specific_heat': PropsSI('C', 'P', pressure, 'T', temperature, self.fluid),
                'enthalpy': PropsSI('H', 'P', pressure, 'T', temperature, self.fluid)
            }
            
            return props
            
        except Exception as e:
            raise ValueError(f"物性计算错误: {str(e)}")
            
    def get_heat_transfer_parameters(self, pressure: float, temperature: float,
                                   velocity: float, hydraulic_diameter: float,
                                   state: FluidState) -> Dict[str, float]:
        """
        计算传热相关参数
        
        参数:
            pressure: 压力 (Pa)
            temperature: 温度 (K)
            velocity: 流速 (m/s)
            hydraulic_diameter: 水力直径 (m)
            state: 流体状态
            
        返回:
            包含以下参数的字典:
            - reynolds: 雷诺数
            - prandtl: 普朗特数
            - nusselt: 努塞尔数 (使用Dittus-Boelter关联式)
            - heat_transfer_coef: 对流传热系数 (W/m²·K)
        """
        try:
            # 获取基本物性
            props = self.get_properties(pressure, temperature, state)
            
            # 计算无量纲参数
            reynolds = props['density'] * velocity * hydraulic_diameter / \
                      props['viscosity']
            prandtl = props['specific_heat'] * props['viscosity'] / \
                     props['conductivity']
            
            # 使用Dittus-Boelter关联式计算努塞尔数
            # 加热时n=0.4,冷却时n=0.3
            n = 0.4  # 假设加热工况
            nusselt = 0.023 * reynolds**0.8 * prandtl**n
            
            # 计算对流传热系数
            h = nusselt * props['conductivity'] / hydraulic_diameter
            
            return {
                'reynolds': reynolds,
                'prandtl': prandtl,
                'nusselt': nusselt,
                'heat_transfer_coef': h
            }
            
        except Exception as e:
            raise ValueError(f"传热参数计算错误: {str(e)}")
            
    def get_saturation_temperature(self, pressure: float) -> float:
        """
        计算饱和温度
        
        参数:
            pressure: 压力 (Pa)
            
        返回:
            饱和温度 (K)
        """
        return PropsSI('T', 'P', pressure, 'Q', 0, self.fluid)
        
    def get_surface_tension(self, temperature: float) -> float:
        """
        计算表面张力
        
        参数:
            temperature: 温度 (K)
            
        返回:
            表面张力 (N/m)
        """
        return PropsSI('I', 'T', temperature, 'Q', 0, self.fluid)
        
    def _get_saturation_properties(self, pressure: float) -> Dict[str, float]:
        """
        计算饱和状态下的物性
        
        参数:
            pressure: 压力 (Pa)
            
        返回:
            饱和状态物性字典
        """
        # 计算饱和温度
        t_sat = self.get_saturation_temperature(pressure)
        
        # 计算液相物性
        props_l = {
            'density_liquid': PropsSI('D', 'P', pressure, 'Q', 0, self.fluid),
            'viscosity_liquid': PropsSI('V', 'P', pressure, 'Q', 0, self.fluid),
            'conductivity_liquid': PropsSI('L', 'P', pressure, 'Q', 0, self.fluid),
            'specific_heat_liquid': PropsSI('C', 'P', pressure, 'Q', 0, self.fluid),
            'enthalpy_liquid': PropsSI('H', 'P', pressure, 'Q', 0, self.fluid)
        }
        
        # 计算气相物性
        props_v = {
            'density_vapor': PropsSI('D', 'P', pressure, 'Q', 1, self.fluid),
            'viscosity_vapor': PropsSI('V', 'P', pressure, 'Q', 1, self.fluid),
            'conductivity_vapor': PropsSI('L', 'P', pressure, 'Q', 1, self.fluid),
            'specific_heat_vapor': PropsSI('C', 'P', pressure, 'Q', 1, self.fluid),
            'enthalpy_vapor': PropsSI('H', 'P', pressure, 'Q', 1, self.fluid)
        }
        
        # 计算其他参数
        props = {
            'temperature': t_sat,
            'surface_tension': self.get_surface_tension(t_sat),
            'latent_heat': props_v['enthalpy_vapor'] - props_l['enthalpy_liquid']
        }
        
        # 合并所有物性
        return {**props, **props_l, **props_v}
        
    def get_quality(self, pressure: float, enthalpy: float) -> float:
        """
        计算干度
        
        参数:
            pressure: 压力 (Pa)
            enthalpy: 比焓 (J/kg)
            
        返回:
            干度 (0~1)
        """
        # 获取饱和状态物性
        props = self._get_saturation_properties(pressure)
        
        # 计算干度
        h_f = props['enthalpy_liquid']
        h_fg = props['latent_heat']
        
        return max(0.0, min(1.0, (enthalpy - h_f) / h_fg))
        
    def get_void_fraction(self, pressure: float, quality: float,
                         alpha_model: str = 'homogeneous') -> float:
        """
        计算空泡分数
        
        参数:
            pressure: 压力 (Pa)
            quality: 干度
            alpha_model: 空泡分数模型 ('homogeneous'/'slip')
            
        返回:
            空泡分数 (0~1)
        """
        props = self._get_saturation_properties(pressure)
        
        if alpha_model == 'homogeneous':
            # 均相流模型
            return quality * props['density_liquid'] / \
                   (quality * props['density_liquid'] + 
                    (1 - quality) * props['density_vapor'])
        else:
            # 滑移流模型 (使用Zivi相关式)
            S = (props['density_vapor'] / props['density_liquid'])**(1/3)
            return (1 + (1-quality)/quality * props['density_vapor']/ \
                   props['density_liquid'] * S)**(-1) 

    def get_temperature(self, pressure: float, enthalpy: float, 
                       phase: str) -> float:
        """
        根据压力和焓计算温度
        
        参数:
            pressure: 压力 (Pa)
            enthalpy: 比焓 (J/kg)
            phase: 相态 ('liquid' 或 'vapor')
            
        返回:
            温度 (K)
        """
        try:
            if phase == 'liquid':
                # 对于液相，使用压力和焓计算温度
                return PropsSI('T', 'P', pressure, 'H', enthalpy, self.fluid)
            elif phase == 'vapor':
                # 对于气相，使用压力和焓计算温度
                return PropsSI('T', 'P', pressure, 'H', enthalpy, self.fluid)
            else:
                raise ValueError(f"未知的相态: {phase}")
        except Exception as e:
            raise ValueError(f"温度计算错误: {str(e)}")